Skip to content

[Pytorch] Enable TE Op to consume extra_outputs from a previously run Op in TE Sequential - #3320

Open
vthumbe1503 wants to merge 39 commits into
NVIDIA:mainfrom
vthumbe1503:enable_extra_out_consumption
Open

[Pytorch] Enable TE Op to consume extra_outputs from a previously run Op in TE Sequential#3320
vthumbe1503 wants to merge 39 commits into
NVIDIA:mainfrom
vthumbe1503:enable_extra_out_consumption

Conversation

@vthumbe1503

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

vthumbe1503 and others added 11 commits July 28, 2026 23:05
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
…h error handling tests

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
vthumbe1503 marked this pull request as ready for review August 5, 2026 23:58
@vthumbe1503
vthumbe1503 requested a review from timmoon10 as a code owner August 5, 2026 23:58
@vthumbe1503
vthumbe1503 requested a review from ptrendx August 6, 2026 00:00
@vthumbe1503 vthumbe1503 changed the title Enable TE Sequential Op to consume extra_outputs from a previously run Op [Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op Aug 6, 2026
@greptile-apps

greptile-apps Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR adds named extra-tensor channels that route outputs from earlier fusible operations into later operations within one OperationFuser.

  • Adds channel configuration and stale-routing detection to basic operations and fusers.
  • Routes channel tensors and accumulated gradients across forward and backward fusion boundaries.
  • Supports public or internal-only channel outputs and validates channel topology.
  • Adds documentation and tests for residual routing, fan-out, fusion, gradients, and invalid configurations.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; current HEAD rejects mutations against retained routing snapshots while temporary and discarded fusers leave no persistent channel locks.

Important Files Changed

Filename Overview
transformer_engine/pytorch/ops/fuser.py Implements channel topology, forward routing, backward gradient accumulation, public slot handling, and stale-configuration rejection; the previously reported lifecycle issues are resolved.
transformer_engine/pytorch/ops/op.py Adds validated channel setters and monotonically versioned channel configuration without retaining fuser ownership or locks.
transformer_engine/pytorch/ops/basic/grouped_linear.py Initializes the dynamic extra-input count before the base class allocates channel metadata.
tests/pytorch/test_fusible_ops.py Adds broad channel-routing coverage across public and hidden outputs, fan-out, gradients, fusion boundaries, validation, and stale configuration.
docs/examples/op_fuser/op_fuser.rst Documents named channel semantics, scope, caller-visible outputs, fusion responsibilities, and configuration lifetime.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Caller["Caller inputs"] --> Fuser["OperationFuser"]
  Producer["Producer basic op"] -->|named extra output| Channel["Fuser-local channel"]
  Channel -->|named extra input| Consumer["Later consumer op"]
  Producer -->|optional public output| CallerOut["Caller outputs"]
  Consumer --> MainOut["Main output"]
  MainOut --> Backward["Backward pass"]
  Backward -->|consumer gradient| ChannelGrad["Accumulated channel gradient"]
  ChannelGrad --> Producer
Loading

Reviews (21): Last reviewed commit: "cleanup" | Re-trigger Greptile

Comment thread transformer_engine/pytorch/ops/op.py
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated

@ptrendx ptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Additional review comments from Codex:

  1. [High] Internal channel outputs lose the signal that their gradient is required.
     transformer_engine/pytorch/ops/fuser.py:170 only calls requires_grad_ for public outputs at line 194. A fresh tensor created by an internal producer inside
     torch.autograd.Function.forward therefore arrives at its consumer with requires_grad=False. Existing operations such as transformer_engine/pytorch/ops/basic/
     swiglu.py:468 and ScaledSReLU use that flag to decide whether to compute the extra-input gradient. They consequently return None, silently dropping the
     gradient to a differentiable producer such as the documented router-probability dispatch. The new tests do not expose this because MakeExtraOutput returns the
     original input tensor.

  2. [Medium] Channel routing violates the declared iterable output contract.
     transformer_engine/pytorch/ops/fuser.py:140 applies len() and indexing to a producer’s extra outputs, but transformer_engine/pytorch/ops/op.py:92 permits any
     Iterable[Iterable[Tensor]]. A custom Dispatch returning a generator works under the old flattening logic but now fails when a later channel consumer executes.

  3. [Medium] A standalone operation call permanently prevents later channel configuration.
     transformer_engine/pytorch/ops/op.py:598 constructs a temporary OperationFuser, while transformer_engine/pytorch/ops/fuser.py:509 permanently locks every
     attached operation. Calling an operation once through its normal forward, then placing it into a channel-connected Sequential, makes either setter raise even
     though the temporary fuser no longer exists.

  4. [Medium] The tests do not exercise two major routing branches.
     tests/pytorch/test_fusible_ops.py:460 registers only a forward fusion, so backward remains unfused and never exercises the same-fusion skip at
     transformer_engine/pytorch/ops/fuser.py:323. The multi-output test at tests/pytorch/test_fusible_ops.py:624 only checks duplicate-name rejection; its custom
     operation never runs. Thus mixed bound/unbound slot ordering, filtered autograd returns, and the modified two-input GroupedLinear(scale_bias=True) behavior
     remain unproved.

  ## Suggested repairs

  - For finding 1: Preserve the gradient-requirement flag on every extra output before classifying it as public or internal, or carry equivalent explicit per-slot
    metadata. Add a producer that creates a fresh tensor and verify gradient propagation through ScaledSwiGLU or ScaledSReLU.

  - For finding 2: Materialize and validate each operation’s extra outputs as a tuple immediately after fuser_forward; store that tuple for later consumers and
    lifetime tracking.

  - For finding 3: Make transient fusers created by BasicOperation.forward non-locking, while persistent Sequential fusers retain immutable routing. Add a call-
    then-bind regression test.

  - For finding 4: Add a joint/backward fused residual operation and assert fusion selection plus input/parameter/channel gradients. Add a successful multi-input/
    multi-output routing test and a GroupedLinear(scale_bias=True) channel case.

I would like you to also take a look at the tests - multiple of them duplicate each other (e.g. test_channel_fan_out_accumulates_grad is a stronger duplicate of test_internal_extra_tensor_channel_fanout).

Comment thread docs/examples/op_fuser/op_fuser.rst Outdated
Comment thread docs/examples/op_fuser/op_fuser.rst Outdated
Comment thread docs/examples/op_fuser/op_fuser.rst Outdated
and cycles are not supported.
- A channel has exactly one producer, but its output may fan out to
multiple consumers.
- Every named output channel must have at least one consumer, and the

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This limitation that there has to be at least one consumer in the named channel seems
arbitrary to me. If we do not strictly need this behavior then we shouldn't have that
as it would introduce friction when somebody needs to refactor the code using those
named channels by splitting the sequential - now they also need to remove the channel
names. In fact, I would expect people to generally want to name their extra outputs and
inputs even if they would not be reused inside the sequential. That could also enable
us to accept and return the dictionary rather than a list (which would make it less
fragile).

@vthumbe1503 vthumbe1503 Aug 10, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Originally I wanted to restrict the channel naming as a way to just do internal routing of tensors to reduce possibility of errors and the friction was kind of intentional. But I see your point of making it more seamless for user in future to construct a big a sequential op. If they want to refactor a code from 1 to 2 below

  1. single sequential having internal routing
  2. Two sequentials with one sequential passing extra output as extra input to another sequential

This can indeed be a problem since there might be some use-cases just supporting 2 but not 1.

And so I have removed that restriction. However, supporting extra_input and extra_output as dictionaries would be a problem from backwards compatibility perspective. Also, I want to restrict the scope of this PR. And allowing for dict based extra_input and extra output can be a seperate PR.

I have one extra requirement from named extra input channel added currently. If two different extra inputs share the same channel name, and is not internally connected to extra output of a previous op. Caller/User should still provide the extra_input two times.

This is done so that user's code doesnt have to change while naming an input channel vs not naming it. Also as you can see introducing extra_input dict is also going to make this tricky from backward compatibility perspective.

@vthumbe1503 vthumbe1503 Aug 11, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update: Tim has a valid point below of extra_input's complexity. For the named extra_input which is not connected to any producer op, we are simply raising an error. Its users responsibility to unname the channel if it is already named(by setting it to None) and is not connected to any internal producer.

vthumbe1503 and others added 5 commits August 7, 2026 22:53
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
vthumbe1503 and others added 11 commits August 9, 2026 22:15
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
…/TransformerEngine into enable_extra_out_consumption
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503 vthumbe1503 changed the title [Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op [Pytorch] Enable TE Op to consume extra_outputs from a previously run Op in TE Sequential Aug 9, 2026
@vthumbe1503

Copy link
Copy Markdown
Collaborator Author

/te-ci L1 pytorch

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Comment thread transformer_engine/pytorch/ops/basic/grouped_linear.py
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated
Comment thread transformer_engine/pytorch/ops/fuser.py Outdated
Comment thread tests/pytorch/test_fusible_ops.py
Comment thread tests/pytorch/test_fusible_ops.py Outdated
Comment thread tests/pytorch/test_fusible_ops.py Outdated
Comment on lines +235 to +238
Channels cannot connect operations in different ``OperationFuser``
instances. In particular, an ordinary PyTorch module inside a
``Sequential`` splits the fusible operations on either side into
separate fusers. The following channel connection is therefore not

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be nice if Sequential could handle channels across OperationFusers, but the implementation would be quite hairy and not worth it for the current effort.

Comment thread docs/examples/op_fuser/op_fuser.rst Outdated
vthumbe1503 and others added 5 commits August 10, 2026 22:07
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
…tted

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@ptrendx ptrendx self-assigned this Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants